
using System;
using TomanuExtensions.Utils;
using System.Diagnostics;
using System.Collections.Generic;
using System.Drawing;
using RaytracerLib.MathLib;
using YAXLib;

namespace RaytracerLib.Samplers
{
    // Taken from GSL 1.9
    public class SobolSampler : LowDiscrepancySequenceSampler
    {
        #region Consts

        private const int MAX_DIMENSION = 40;

        // bit count; assumes sizeof(int) >= 32-bit
        private const int BIT_COUNT = 30;

        // primitive polynomials in binary encoding
        private readonly static int[] s_primitive_polynomials =
        {
            1,     3,   7,  11,  13,  19,  25,  37,  59,  47,
            61,   55,  41,  67,  97,  91, 109, 103, 115, 131, 
            193, 137, 145, 143, 241, 157, 185, 167, 229, 171,
            213, 191, 253, 203, 211, 239, 247, 285, 369, 299
        };

        // initial values for direction tables, following
        // Bratley+Fox, taken from [Sobol+Levitan, preprint 1976]
        private readonly static int[,] s_dir_init = 
        {
            {
                0, 1, 1, 1, 1, 1, 1, 1, 1, 1,
                1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
                1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
                1, 1, 1, 1, 1, 1, 1, 1, 1, 1
            },
            {
                0, 0, 1, 3, 1, 3, 1, 3, 3, 1,
                3, 1, 3, 1, 3, 1, 1, 3, 1, 3,
                1, 3, 1, 3, 3, 1, 3, 1, 3, 1,
                3, 1, 1, 3, 1, 3, 1, 3, 1, 3
            }, 
            {
                0, 0, 0, 7, 5, 1, 3, 3, 7, 5,
                5, 7, 7, 1, 3, 3, 7, 5, 1, 1,
                5, 3, 3, 1, 7, 5, 1, 3, 3, 7,
                5, 1, 1, 5, 7, 7, 5, 1, 3, 3
            }, 
            {
                0,  0,  0,  0,  0,  1,  7,  9, 13, 11,
                1,  3,  7,  9,  5, 13, 13, 11,  3, 15,
                5,  3, 15,  7,  9, 13,  9,  1, 11,  7,
                5, 15,  1, 15, 11,  5,  3,  1,  7,  9
            }, 
            {
                0,   0,  0,  0,  0,  0,  0,  9,  3, 27,
                15, 29, 21, 23, 19, 11, 25,  7, 13, 17,
                1,  25, 29,  3, 31, 11,  5, 23, 27, 19,
                21,  5,  1, 17, 13,  7, 15,  9, 31,  9
            }, 
            {
                    0,  0,  0,  0,  0,  0,  0,  0,  0,  0,
                    0,  0,  0, 37, 33,  7,  5, 11, 39, 63,
                27, 17, 15, 23, 29,  3, 21, 13, 31, 25,
                    9, 49, 33, 19, 29, 11, 19, 27, 15, 25
            }, 
            {
                    0,   0,  0,  0,  0,  0,    0,  0,  0,   0,
                    0,   0,  0,  0,  0,  0,    0,  0,  0,  13,
                33, 115, 41, 79, 17,  29, 119, 75, 73, 105,
                    7,  59, 65, 21,  3, 113,  61, 89, 45, 107
            }, 
            {
                0, 0, 0, 0, 0, 0, 0, 0,  0,  0,
                0, 0, 0, 0, 0, 0, 0, 0,  0,  0,
                0, 0, 0, 0, 0, 0, 0, 0,  0,  0,
                0, 0, 0, 0, 0, 0, 0, 7, 23, 39
            }
        };

        #endregion

        // number of calls with this generator
        private uint m_n = 0;

        // 1/denominator for last numerator vector
        private double m_last_denominator_inv;

        // last generated numerator vector
        private int[] m_last_numerator_vec;

        // direction number table
        private int[,] m_v_direction;

        public uint Dimension = 2;

        private void Init()
        {
            if (Dimension < 1 || Dimension > MAX_DIMENSION)
                throw new NotImplementedException();

            m_last_numerator_vec = new int[Dimension];
            m_v_direction = new int[BIT_COUNT, Dimension];

            // Initialize direction table in dimension 0. 
            for (int k = 0; k < BIT_COUNT; k++)
                m_v_direction[k, 0] = 1;

            int ell = 0;

            // Initialize in remaining dimensions.
            for (int dim = 1; dim < Dimension; dim++)
            {
                int poly_index = dim;

                int p_i = s_primitive_polynomials[poly_index];

                int poly_degree = Bits.Highest(p_i) - 1;
                int[] includ = new int[poly_degree];

                // Expand the polynomial bit pattern to separate
                // components of the logical array includ[].
                for (int k = poly_degree - 1; k >= 0; k--)
                {
                    includ[k] = p_i % 2;
                    p_i >>= 2;
                }

                // Leading elements for dimension i come from m_v_direction[][].
                for (int j = 0; j < poly_degree; j++)
                    m_v_direction[j, dim] = s_dir_init[j, dim];

                // Calculate remaining elements for this dimension,
                // as explained in Bratley+Fox, section 2.
                for (int j = poly_degree; j < BIT_COUNT; j++)
                {
                    int newv = m_v_direction[j - poly_degree, dim];
                    ell = 1;
                    for (int k = 0; k < poly_degree; k++)
                    {
                        ell *= 2;
                        if (includ[k] == 1)
                            newv ^= (ell * m_v_direction[j - k - 1, dim]);
                    }
                    m_v_direction[j, dim] = newv;
                }
            }

            // Multiply columns of v by appropriate power of 2.
            ell = 1;
            for (int j = BIT_COUNT - 1 - 1; j >= 0; j--)
            {
                ell *= 2;
                for (int dim = 0; dim < Dimension; dim++)
                    m_v_direction[j, dim] *= ell;
            }

            // 1/(common denominator of the elements in v_direction) 
            m_last_denominator_inv = 1.0 / (2.0 * ell);
        }

        private void NextSample(double[] a_v)
        {
            Debug.Assert(a_v.Length == Dimension);

            int bitssets = Bits.LeastSignificantZero(m_n) + 1;

            // Check for exhaustion.
            if (bitssets > BIT_COUNT)
                throw new NotImplementedException();

            for (int dim = 0; dim < Dimension; dim++)
            {
                int direction_i = m_v_direction[bitssets - 1, dim];
                int old_numerator_i = m_last_numerator_vec[dim];
                int new_numerator_i = old_numerator_i ^ direction_i;
                m_last_numerator_vec[dim] = new_numerator_i;
                a_v[dim] = new_numerator_i * m_last_denominator_inv;
            }

            m_n++;
        }

        protected override Vector2 NextSample()
        {
            double[] v = new double[Dimension];
            NextSample(v);
            return new Vector2(v[0], v[1]);
        }

        internal override void RenderStart(RenderStartPhase a_phase)
        {
            base.RenderStart(a_phase);

            if (a_phase == RenderStartPhase.PrepareObjectToRender)
                Init();
        }

        public override SamplerType SamplerType
        {
            get 
            {
                return SamplerType.Sobol;
            }
        }
    }
}